diff --git a/model.py b/model.py
index 0134c39..0356ad5 100755
--- a/model.py
+++ b/model.py
@@ -8,7 +8,10 @@ from torch import nn
 from torch.nn import functional as F
 from torch.autograd import Function

-from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d, conv2d_gradfix
+from models.networks.op import fused_leaky_relu
+from models.networks.op import FusedLeakyReLU
+from models.networks.op import upfirdn2d
+from models.networks.op import conv2d_gradfix


 class PixelNorm(nn.Module):
@@ -366,27 +369,35 @@ class StyledConv(nn.Module):
         return out


+
 class ToRGB(nn.Module):
-    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
+  """ToRGB."""

-        if upsample:
-            self.upsample = Upsample(blur_kernel)
+  def __init__(self,
+               in_channel,
+               style_dim,
+               output_channel=3,
+               upsample=True,
+               blur_kernel=[1, 3, 3, 1]):
+    super().__init__()

-        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
-        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+    if upsample:
+      self.upsample = Upsample(blur_kernel)

-    def forward(self, input, style, skip=None):
-        out = self.conv(input, style)
-        out = out + self.bias
+    self.conv = ModulatedConv2d(
+        in_channel, output_channel, 1, style_dim, demodulate=False)
+    self.bias = nn.Parameter(torch.zeros(1, output_channel, 1, 1))

-        if skip is not None:
-            skip = self.upsample(skip)
+  def forward(self, input_, style, skip=None):
+    out = self.conv(input_, style)
+    out = out + self.bias

-            out = out + skip
+    if skip is not None:
+      skip = self.upsample(skip)

-        return out
+      out = out + skip

+    return out

 class Generator(nn.Module):
     def __init__(
